import pygame
import math
import random

# Initialize Pygame
pygame.init()

# Constants
SCREEN_WIDTH = 800
SCREEN_HEIGHT = 600
GRAVITY = 0.1
THRUST_POWER = 0.2
FUEL_CONSUMPTION = 2
MAX_LANDING_SPEED = 2.0
MAX_LANDING_ANGLE = 15

# Colors
BLACK = (0, 0, 0)
WHITE = (255, 255, 255)
RED = (255, 0, 0)
GREEN = (0, 255, 0)
YELLOW = (255, 255, 0)
BLUE = (0, 0, 255)
GRAY = (128, 128, 128)
ORANGE = (255, 165, 0)


class LunarLander:
    def __init__(self, x, y):
        self.x = x
        self.y = y
        self.vel_x = 0
        self.vel_y = 0
        self.angle = 0
        self.fuel = 500
        self.thrust = False
        self.width = 20
        self.height = 30
        self.landed = False
        self.crashed = False

    def update(self, keys):
        # Handle input
        self.thrust = False

        if keys[pygame.K_LEFT] and self.fuel > 0:
            self.angle -= 3
        if keys[pygame.K_RIGHT] and self.fuel > 0:
            self.angle += 3
        if keys[pygame.K_UP] and self.fuel > 0:
            self.thrust = True
            # Apply thrust in the direction the lander is pointing
            thrust_x = math.sin(math.radians(self.angle)) * THRUST_POWER
            thrust_y = -math.cos(math.radians(self.angle)) * THRUST_POWER
            self.vel_x += thrust_x
            self.vel_y += thrust_y
            self.fuel -= FUEL_CONSUMPTION

        # Apply gravity
        self.vel_y += GRAVITY

        # Update position
        self.x += self.vel_x
        self.y += self.vel_y

        # Wrap around screen horizontally
        if self.x < 0:
            self.x = SCREEN_WIDTH
        elif self.x > SCREEN_WIDTH:
            self.x = 0

        # Keep angle in reasonable range
        self.angle = self.angle % 360

    def draw(self, screen):
        # Calculate lander corners based on angle
        cos_a = math.cos(math.radians(self.angle))
        sin_a = math.sin(math.radians(self.angle))

        # Lander body points (relative to center)
        points = [
            (-10, -15),  # top left
            (10, -15),  # top right
            (10, 15),  # bottom right
            (-10, 15)  # bottom left
        ]

        # Rotate and translate points
        rotated_points = []
        for px, py in points:
            rx = px * cos_a - py * sin_a + self.x
            ry = px * sin_a + py * cos_a + self.y
            rotated_points.append((rx, ry))

        # Draw lander body
        pygame.draw.polygon(screen, WHITE, rotated_points)

        # Draw thrust flame if thrusting
        if self.thrust and self.fuel > 0:
            flame_length = random.randint(15, 25)
            flame_x = self.x + sin_a * (15 + flame_length // 2)
            flame_y = self.y + cos_a * (15 + flame_length // 2)
            pygame.draw.circle(screen, ORANGE, (int(flame_x), int(flame_y)), 5)
            pygame.draw.circle(screen, YELLOW, (int(flame_x), int(flame_y)), 3)


class Terrain:
    def __init__(self):
        self.points = []
        self.landing_zones = []
        self.generate_terrain()

    def generate_terrain(self):
        # Generate random terrain
        num_points = 20
        for i in range(num_points + 1):
            x = (SCREEN_WIDTH / num_points) * i
            # Create some flat landing zones
            if i in [5, 10, 15]:  # Landing zones
                y = SCREEN_HEIGHT - random.randint(100, 200)
                # Make flat landing zone
                if i < num_points:
                    self.points.append((x, y))
                    next_x = (SCREEN_WIDTH / num_points) * (i + 1)
                    self.points.append((next_x, y))
                    self.landing_zones.append((x, next_x, y))
                    i += 1
            else:
                y = SCREEN_HEIGHT - random.randint(50, 300)
                self.points.append((x, y))

    def draw(self, screen):
        # Draw terrain
        if len(self.points) > 1:
            pygame.draw.lines(screen, WHITE, False, self.points, 2)

        # Draw landing zones in green
        for start_x, end_x, y in self.landing_zones:
            pygame.draw.line(screen, GREEN, (start_x, y), (end_x, y), 4)

    def check_collision(self, lander):
        # Simple collision detection
        lander_bottom = lander.y + lander.height // 2

        # Find terrain height at lander position
        for i in range(len(self.points) - 1):
            x1, y1 = self.points[i]
            x2, y2 = self.points[i + 1]

            if x1 <= lander.x <= x2:
                # Linear interpolation for terrain height
                ratio = (lander.x - x1) / (x2 - x1) if x2 != x1 else 0
                terrain_y = y1 + (y2 - y1) * ratio

                if lander_bottom >= terrain_y:
                    # Check if it's a safe landing
                    speed = math.sqrt(lander.vel_x ** 2 + lander.vel_y ** 2)
                    angle_ok = abs(lander.angle) <= MAX_LANDING_ANGLE or abs(lander.angle - 360) <= MAX_LANDING_ANGLE

                    # Check if landed on a landing zone
                    on_landing_zone = any(start_x <= lander.x <= end_x and abs(terrain_y - y) < 10
                                          for start_x, end_x, y in self.landing_zones)

                    if speed <= MAX_LANDING_SPEED and angle_ok and on_landing_zone:
                        lander.landed = True
                    else:
                        lander.crashed = True
                    return True
        return False


class Game:
    def __init__(self):
        self.screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT))
        pygame.display.set_caption("Lunar Lander")
        self.clock = pygame.time.Clock()
        self.font = pygame.font.Font(None, 36)
        self.small_font = pygame.font.Font(None, 24)
        self.reset_game()

    def reset_game(self):
        self.lander = LunarLander(SCREEN_WIDTH // 2, 50)
        self.terrain = Terrain()
        self.game_over = False
        self.message = ""

    def draw_hud(self):
        # Fuel gauge
        fuel_text = self.small_font.render(f"Fuel: {max(0, int(self.lander.fuel))}", True, WHITE)
        self.screen.blit(fuel_text, (10, 10))

        # Speed indicator
        speed = math.sqrt(self.lander.vel_x ** 2 + self.lander.vel_y ** 2)
        speed_text = self.small_font.render(f"Speed: {speed:.1f}", True, WHITE)
        self.screen.blit(speed_text, (10, 35))

        # Angle indicator
        display_angle = self.lander.angle if self.lander.angle <= 180 else self.lander.angle - 360
        angle_text = self.small_font.render(f"Angle: {display_angle:.0f}°", True, WHITE)
        self.screen.blit(angle_text, (10, 60))

        # Instructions
        if not self.game_over:
            instructions = [
                "Arrow Keys: Rotate and Thrust",
                "Land gently on green landing zones",
                f"Max speed: {MAX_LANDING_SPEED}, Max angle: ±{MAX_LANDING_ANGLE}°"
            ]
            for i, instruction in enumerate(instructions):
                text = self.small_font.render(instruction, True, WHITE)
                self.screen.blit(text, (SCREEN_WIDTH - 350, 10 + i * 20))

    def run(self):
        running = True

        while running:
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    running = False
                elif event.type == pygame.KEYDOWN:
                    if event.key == pygame.K_r and self.game_over:
                        self.reset_game()

            if not self.game_over:
                keys = pygame.key.get_pressed()
                self.lander.update(keys)

                # Check for collision with terrain
                if self.terrain.check_collision(self.lander):
                    self.game_over = True
                    if self.lander.landed:
                        self.message = "SUCCESSFUL LANDING!"
                    else:
                        self.message = "CRASHED!"

                # Check if lander went off screen (top)
                if self.lander.y < -50:
                    self.lander.y = SCREEN_HEIGHT + 50

                # Check if out of fuel and falling
                if self.lander.fuel <= 0 and self.lander.vel_y > 5:
                    self.message = "OUT OF FUEL!"

            # Draw everything
            self.screen.fill(BLACK)
            self.terrain.draw(self.screen)
            self.lander.draw(self.screen)
            self.draw_hud()

            # Draw game over message
            if self.game_over:
                if self.lander.landed:
                    color = GREEN
                else:
                    color = RED

                text = self.font.render(self.message, True, color)
                text_rect = text.get_rect(center=(SCREEN_WIDTH // 2, SCREEN_HEIGHT // 2))
                self.screen.blit(text, text_rect)

                restart_text = self.small_font.render("Press R to restart", True, WHITE)
                restart_rect = restart_text.get_rect(center=(SCREEN_WIDTH // 2, SCREEN_HEIGHT // 2 + 50))
                self.screen.blit(restart_text, restart_rect)

            pygame.display.flip()
            self.clock.tick(60)

        pygame.quit()


# Run the game
if __name__ == "__main__":
    game = Game()
    game.run()